FROM nvcr.io/nvidia/pytorch:20.12-py3

RUN pip install pillow
RUN pip install cython
RUN pip install matplotlib
RUN pip install tensorboard
RUN pip install tensorboardX
RUN pip install tfrecord
RUN pip install git+https://github.com/rtqichen/torchdiffeq
RUN apt update
RUN apt install screen -y
